# =============================================================================
# =============================================================================
# Plot a selection cohorts. A variation of this code was used to generate Appendix Figure E2
# =============================================================================
# =============================================================================

#Geographic layer of analysis
geo_year = 2019
cbgf = pd.read_parquet(cdd['p_d_geo'] + r'\cbg_geo'+str(geo_year)+'.parquet').rename(columns={'cbg_lat':'clat','cbg_long':'clong'})
cbg = cbgf[['fips','cbg','clat','clong']].copy()


# =============================================================================
#Import the cohorts
tcp = pd.read_parquet(cdd['p_d_geo_tcp'] + r'\\tcp_25k_'+ts['geog']+'_'+ts['freq']+'.parquet')

# Subset the tcp panel to the desired radius around treated zones.
tcps = tcp[(tcp.distance<10000)|((tcp.treated==1)&(tcp.distance<10000))|((tcp.ctype==0)&(tcp.distance_t2cf<25000))].copy()
tcps = tcps[tcps['open'] > dt.datetime(2014,6,30)]

#Keep only the 20 closest counterfactuals or the 20 lowest ranked counterfactuals.
tcps=tcps[(tcps.groupby('cohort')['treated'].transform('max')!=tcps.groupby('cohort')['treated'].transform('min'))]

tcps = tcps[tcps.etime==0].drop_duplicates(subset=['cohort','cohort_cf',ts['geog'],'ctype']).copy()
tcps['tract'] = np.floor(tcps['cbg']/10)
tcps['distance_t2cf'] = tcps['distance_t2cf'].fillna(0)
tcps['open_cf'] = tcps['open_cf'].fillna(dt.datetime(2000,1,1))
tcps['ctype'] = tcps['ctype'].map({0:3,1:1,2:2,5:0})
tcps['ctypes'] = tcps['ctype'].map({3:'local',1:'future',2:'past',0:'treated'})


#Import tract controls to merge in some info about the location
tract_controls = pd.read_parquet(cdd['p_d_acs'] + r'\tract\acs_2010_2021_s3.parquet')
tract_controls = tract_controls[tract_controls.year==2018]

#Drop any duplicate observations and then merge with geog.
tcp_geo = pd.merge(tcps,cbgf,how='left',on=ts['geog'])
tcp_geo = pd.merge(tcp_geo,tract_controls,how='left',on='tract')

#Read in plasma center list.
pcp = pd.read_parquet(cdd['p_d_pcp'] + r'\pcp.parquet')
pcp['color'] = pcp['open'].dt.year.apply(lambda x: 'darkred' if x<=2014 else 'peru')
pcp['name'] = pcp['name_ref'].fillna(pcp['name_elgl']).fillna(pcp['name_app']).fillna('other')



# =============================================================================
# Make custom plot for the Paper
# =============================================================================

# =============================================================================
# Read in treatment intensity to color all of the other (non-treated regions)
treat = pd.read_parquet(cdd['p_d_geo_tcp'] + r'\treat_'+str(int(ts['radius']/1000))+'k_'+ts['geog']+'_'+ts['freq']+'.parquet')
treat = pd.merge(treat,cbgf[[ts['freq'],'clat','clong','geometry']],how='left',on=ts['freq'])


dimensions = {'width':3200,'height':1800}
cohl = {
        844154354091:{'place':'Houston, TX',"lat":1000,"long":1000,"zoom":10,'linew':3,'pointw':10,'show_legend':False}
        }



tgt = list(cohl.keys())[0]
for tgt in list(cohl.keys()):
    tcp_geo_tgt = tcp_geo[tcp_geo.cohort==tgt].copy()

    #Make the tgt a geodataframe (or choropleth won't work)
    tcp_geo_tgt = gpd.GeoDataFrame(tcp_geo_tgt.copy(), geometry=tcp_geo_tgt['geometry'])
    tcp_geo_tgt.crs = {'init' :'epsg:4269'}
    tcp_geo_tgt = tcp_geo_tgt.to_crs(epsg=4269)
    # tcp_geo_tgt.crs
    tcp_geo_tgt['sqkm'] = tcp_geo_tgt['ALAND'] * 0.09290304 / 1000000
    
    coh_open = tcp_geo_tgt['open'].iloc[0]
    coh_open_dt = tcp_geo_tgt['date'].iloc[0]
    coh_openy = pcp[pcp.mfeiu==tgt]['open'].dt.year.iloc[0]
    coh_city = pcp[pcp.mfeiu==tgt]['city'].iloc[0]
    coh_st = pcp[pcp.mfeiu==tgt]['state'].iloc[0]
    coh_clat = pcp[pcp.mfeiu==tgt]['clat'].iloc[0]
    coh_clong = pcp[pcp.mfeiu==tgt]['clong'].iloc[0]
    coh_clat_fc = pcp[pcp.mfeiu.isin(tcp_geo_tgt.cohort_cf.unique())]['clat'].mean()
    coh_clong_fc = pcp[pcp.mfeiu.isin(tcp_geo_tgt.cohort_cf.unique())]['clong'].mean()
    
    DemoGraphs = {'pvf100_p_wq10':{'name':'Poverty (Pct. Pop. <100FPL  Decile)','color':"OrRd"},
                   'educ_bach_p_wq10':{'name':'Education (Pct. Pop. w/ Bach Decile)','color':"BuGn"},
                   'pop2sqkm_wq10':{'name':'Pop Debsity (Pop. per SqKm  Decile)','color':"BuGn"}}
    
    
    #Construct a graph
    #   https://georgetsilva.github.io/posts/mapping-points-with-folium/
    #   https://stackoverflow.com/questions/52428916/python-folium-markercluster-color-customization
    #   SEE THIS FOR USING GEOJSON TO CHOROPLETH with tooltips https://vverde.github.io/blob/interactivechoropleth.html
    #   or this https://towardsdatascience.com/how-to-step-up-your-folium-choropleth-map-skills-17cf6de7c6fe
    #   List of tiles https://deparkes.co.uk/2016/06/10/folium-map-tiles/   (I like stamentoner, and cartodbpositron)
    coh_lat_g = cohl[tgt]['lat'] if cohl[tgt]['lat'] < 1000 else coh_clat_fc
    coh_long_g = cohl[tgt]['long'] if cohl[tgt]['long'] < 1000 else coh_clong_fc
    
    m = folium.Map(location=[coh_lat_g, coh_long_g], tiles='cartodbpositron',zoom_start=cohl[tgt]['zoom'])
    
    #Each degree of lat or long is roughly 40 miles or 80 km so to get a very rough search grid around a point of radius r (in km) I need the degree difference to be less than r/100
    radius=50000
    degdist = 2*radius/100000
    treat_tgt = treat[(treat.date==coh_open_dt)&(abs(treat.clat-coh_clat)<degdist)&(abs(treat.clong-coh_clong)<degdist)].copy()
    treat_tgt['distance']=treat_tgt.apply(lambda x: vinc.vincenty_inverse((coh_clat,coh_clong),(x['clat'],x['clong'])).m,axis=1)
    
    treat_tgt = gpd.GeoDataFrame(treat_tgt.copy(), geometry=treat_tgt['geometry'])
    treat_tgt.crs = {'init' :'epsg:4269'}
    treat_tgt = treat_tgt.to_crs(epsg=4269)
    treat_tgt_s = treat_tgt[['closest_cohort','geometry']].dissolve(by='closest_cohort').reset_index()
    treat_tgt_s['ctype'] = np.nan
    
    # Here we add cross-hatching (crossing lines) to areas which we want to study in the main sample. 
    #   I don't like the built in cross-hatching so I am sticking with a boundary instead.
    tcp_geo_tgt_ch = tcp_geo_tgt[(tcp_geo_tgt['distance']<5000)&(tcp_geo_tgt['intensity_absdelta_tcf']>5000)][['cohort_cf','geometry']].dissolve(by='cohort_cf').reset_index()
    tcp_geo_tgt_ch['ctype'] = np.nan
    choro_hc = folium.Choropleth(
        geo_data=tcp_geo_tgt_ch[['cohort_cf', 'ctype', 'geometry']],
        name="Dist < 5km, IC > 5km",
        data=tcp_geo_tgt_ch[['cohort_cf', 'ctype', 'geometry']],
        columns=['cohort_cf', 'ctype'],
        key_on="feature.properties.cohort_cf",
        fill_color="PuBuGn",
        nan_fill_opacity=0,
        nan_fill_color='black',
        fill_opacity=0,
        line_opacity=1,
        line_color='black',
        line_weight=cohl[tgt]['linew']*2,
        legend_name="Cohort Border",
        highlight=True,
        show=True)
    for key in choro_hc._children:
        if key.startswith('color_map'):
            del(choro_hc._children[key])
    choro_hc.add_to(m)
    
    pcp2 = pcp.copy()    
    radius=100000
    degdist = 2*radius/100000
    pcp2 = pcp2[(abs(pcp2.clat-coh_clat_fc)<degdist)&(abs(pcp2.clong-coh_clong_fc)<degdist)].copy()

    pcp2['color'] = pcp2.apply(lambda x: 'blue' if x['open']<=coh_open else 'purple',axis=1)
    pcp2['color'] = pcp2.apply(lambda x: 'black' if tgt==x['mfeiu'] else x['color'],axis=1)
    pcp2['hex_color'] = pcp2['color'].apply(lambda x: matplt.colors.to_hex(x))
    pcp2['stype'] = pcp2['color'].map({'blue':'square','purple':'circle','black':'star'})

    pcmg = folium.FeatureGroup(name='PLASMA CENTERS (Toggle All On/Off)',show=True).add_to(m)
    pc_markers_tgt = folium.plugins.FeatureGroupSubGroup(pcmg,name='Plasma Center Target',show=True).add_to(m)
    pc_markers_pastcf = folium.plugins.FeatureGroupSubGroup(pcmg,name='Plasma Centers (past cf)',show=True).add_to(m)
    pc_markers_futurecf = folium.plugins.FeatureGroupSubGroup(pcmg,name='Plasma Centers (future cf)',show=True).add_to(m)
    pc_markers_other = folium.plugins.FeatureGroupSubGroup(pcmg,name='Plasma Centers (other)',show=True).add_to(m)
    pc_markers_old = folium.plugins.FeatureGroupSubGroup(pcmg,name='Plasma Centers (closed <2008)',show=False).add_to(m)
    pcp2['radius'] = cohl[tgt]['pointw'] * (2 + 2 * (pcp2['mfeiu'].isin(tcp_geo_tgt.cohort_cf.unique())) + 0 * (pcp2['mfeiu'].isin(tcp_geo_tgt.cohort.unique())) )
    for pc in pcp2.to_dict(orient='records'):
        tooltip = ("Cohort MFEI: {mfei}<br>"
                  "Name: {name}<br>"
                    "Address: {addr}<br>"
                    "Zip code: {zipc}<br>"
                    "Open Date: {openc}<br>"
                    "Close Date: {closec}<br>"
                  ).format(mfei=str(int(pc['mfeiu'])),
                            addr=str(pc['address']),
                            name=pc['name'],
                            zipc=str(pc['zip']),
                            openc='Future' if pc['open']>dt.datetime(2020,6,1) else str(pc['open'].strftime("%m/%Y")),
                            closec='Open' if pc['open']>dt.datetime(2020,6,1) else str(pc['close'].strftime("%m/%Y"))
                            )

        
        #Define some shapes to use as icons (see https://github.com/masajid390/BeautifyMarker)
        icon = BeautifyIcon(icon=pc['stype'], border_width=10 ,
                                  inner_icon_style='color:'+pc['hex_color']+';font-size:'+str(pc['radius'])+'px;',
                                  border_color='transparent', background_color ='transparent',
                                     text_color=pc['hex_color'])
        marker = folium.Marker(location=(pc['clat'],pc['clong']),
                         tooltip=tooltip,z_index_offset=3,icon=icon)
        
        if (pc['mfeiu'] in (tcp_geo_tgt.cohort.unique()))&(pc['close']>=dt.datetime(2008,1,1)):
            marker.add_to(pc_markers_tgt)
        if (pc['mfeiu'] in (tcp_geo_tgt[tcp_geo_tgt.ctype==1].cohort_cf.unique()))&(pc['close']>=dt.datetime(2008,1,1)):
            marker.add_to(pc_markers_futurecf)
        if (pc['mfeiu'] in (tcp_geo_tgt[tcp_geo_tgt.ctype==2].cohort_cf.unique()))&(pc['close']>=dt.datetime(2008,1,1)):
            marker.add_to(pc_markers_pastcf)
        if (not (pc['mfeiu'] in (list(tcp_geo_tgt.cohort.unique()) + list(tcp_geo_tgt.cohort_cf.unique()))))&(pc['close']>=dt.datetime(2008,1,1)):
            marker.add_to(pc_markers_other)
        # if (pc['close']<dt.datetime(2008,1,1)):
        #     marker.add_to(pc_markers_old)
    
    
    
    # MeasureControl(position='topright', primary_length_unit='meters', secondary_length_unit='miles', primary_area_unit='sqmeters', secondary_area_unit='acres').add_to(m)
    # folium.LayerControl().add_to(m)
    htmlf = cdd['p_rf'] + r'\General\Treatment Cohorts\CBG_cohort_'+str(int(tgt))+'_('+coh_city+', '+coh_st+')_'+str(int(coh_openy))+'.html'
    pngf = cdd['p_rf'] + r'\General\Treatment Cohorts\CBG_cohort_'+str(int(tgt))+'_('+coh_city+', '+coh_st+')_'+str(int(coh_openy))+'.png'
    m.save(htmlf)
    html2png(html=htmlf,png=pngf,driver_exe=cdd['driver_chrome'],dimensions=dimensions)
    webbrowser.open(cdd['p_rf'] + r'\General\Treatment Cohorts\cbg_cohort_'+str(int(tgt))+'_('+coh_city+', '+coh_st+')_'+str(int(coh_openy))+'.html')







